library(coda)
library(locfit)
set.seed(519371)


load("pairData-AttnCheckPassed.Rda")

birthyearData <- read.csv(stringsAsFactors = FALSE,
                          file="./QualtricsRespondentBirthYear-nofraud-AttnCheckPassed.csv")

QualtricsData <- read.csv(stringsAsFactors = FALSE,
                          file="./QualtricOtherQuestionAnswers-nofraud-AttnCheckPassed.csv")


news.outlets <- read.csv(stringsAsFactors = FALSE,
                         file="./news_outlet_political_leaning.csv")

news.outlets$news.leaning.numeric <- NA
news.outlets$news.leaning.numeric[news.outlets$leaning == "Left"] <- -2
news.outlets$news.leaning.numeric[news.outlets$leaning == "Lean Left"] <- -1
news.outlets$news.leaning.numeric[news.outlets$leaning == "Center"] <- 0
news.outlets$news.leaning.numeric[news.outlets$leaning == "Lean Right"] <- 1
news.outlets$news.leaning.numeric[news.outlets$leaning == "Right"] <- 2



load("out-2dDP-nofraud-balanced-newIDs-AttnCheckPassed.Rda")
out <- out.2dDP
gammas <- out[, grep("gamma", colnames(out))]
gamma.hat <- colMeans(gammas)
M <- nrow(out) ## number of MCMC draws





gamma.names <- strsplit(names(gamma.hat), "\\.")
age.reordered <- NULL
educ.reordered <- NULL
race.reordered <- NULL
ideol.reordered <- NULL
mask.reordered <- NULL
socdist.reordered <- NULL
gender.reordered <- NULL
strong.repub.reordered <- NULL
strong.dem.reordered <- NULL
news.lean.reordered <- NULL
timer.reordered <- NULL
party1.reordered <- NULL
party2.reordered <- NULL
party3.reordered <- NULL
party4.reordered <- NULL
party5.reordered <- NULL
party6.reordered <- NULL
for (i in 1:length(gamma.names)){
    rid.i <- gamma.names[[i]][2]

    timer.i <- QualtricsData$Duration..in.seconds.[QualtricsData$rid == rid.i]
    timer.reordered <- c(timer.reordered, timer.i)

    age.i <- QualtricsData$R.age_1[QualtricsData$rid == rid.i]
    age.reordered <- c(age.reordered, age.i)

    educ.i <- QualtricsData$R.edu[QualtricsData$rid == rid.i]
    educ.reordered <- c(educ.reordered, educ.i)

    race.i <- QualtricsData$R.race[QualtricsData$rid == rid.i]
    race.reordered <- c(race.reordered, race.i)    
    
    ideol.i <- QualtricsData$R.ideology[QualtricsData$rid == rid.i]
    ideol.reordered <- c(ideol.reordered, ideol.i)

    mask.i <- QualtricsData$Q1795[QualtricsData$rid == rid.i]
    mask.reordered <- c(mask.reordered, mask.i)

    socdist.i <- QualtricsData$Q1796[QualtricsData$rid == rid.i]
    socdist.reordered <- c(socdist.reordered, socdist.i)

    gender.i <- QualtricsData$R.gender[QualtricsData$rid == rid.i]
    gender.reordered <- c(gender.reordered, gender.i)

    party1.i <- QualtricsData$R.party1[QualtricsData$rid == rid.i]
    party1.reordered <- c(party1.reordered, party1.i)
    
    party2.i <- QualtricsData$R.party2[QualtricsData$rid == rid.i]
    party2.reordered <- c(party2.reordered, party2.i)
    
    party3.i <- QualtricsData$R.party3[QualtricsData$rid == rid.i]
    party3.reordered <- c(party3.reordered, party3.i)
    
    party4.i <- QualtricsData$R.party4[QualtricsData$rid == rid.i]
    party4.reordered <- c(party4.reordered, party4.i)
    
    party5.i <- QualtricsData$R.party5[QualtricsData$rid == rid.i]
    party5.reordered <- c(party5.reordered, party5.i)
    
    party6.i <- QualtricsData$R.party6[QualtricsData$rid == rid.i]
    party6.reordered <- c(party6.reordered, party6.i)
    
    strong.repub.i <- QualtricsData$R.party6[QualtricsData$rid == rid.i]
    if (!is.na(QualtricsData$R.party1[QualtricsData$rid == rid.i])){
        if (is.na(strong.repub.i)){
            strong.repub.i <- 0
        }
    }
    strong.repub.reordered <- c(strong.repub.reordered, strong.repub.i)
    
    strong.dem.i <- QualtricsData$R.party5[QualtricsData$rid == rid.i]
    if (!is.na(QualtricsData$R.party1[QualtricsData$rid == rid.i])){
        if (is.na(strong.dem.i)){
            strong.dem.i <- 0
        }
    }
    strong.dem.reordered <- c(strong.dem.reordered, strong.dem.i)

    news.i <- QualtricsData$R.news[QualtricsData$rid == rid.i]
    news.lean.i <- NA
    if (is.na(news.i) == FALSE){ ## "other responses get NA"
        if (news.i < 26){
            news.lean.i <- news.outlets$news.leaning.numeric[news.outlets$QualtricsCode == news.i]
        }
    }
    news.lean.reordered <- c(news.lean.reordered, news.lean.i)
    
}

nmask <- NULL
for (i in 1:length(mask.reordered)){
    holder <- strsplit(mask.reordered[i], ",")
    nmask <- c(nmask, length(holder[[1]]))
}
nosocdist <- socdist.reordered >= 5
female <- gender.reordered == 2
strong.repub <- strong.repub.reordered == 1
strong.dem <- strong.dem.reordered == 1

mydata <- data.frame(gamma.hat, age=age.reordered, educ=educ.reordered,
                     collegeplus=(educ.reordered>=4),
                     ideology=ideol.reordered, mask=mask.reordered,
                     nmask=nmask, social.distancing=socdist.reordered,
                     nosocdist=nosocdist, female=female,
                     strong.repub=strong.repub, strong.dem=strong.dem,
                     news.lean=news.lean.reordered, race=race.reordered,
                     timer.sec=timer.reordered,
                     party1=party1.reordered,
                     party2=party2.reordered,
                     party3=party3.reordered,
                     party4=party4.reordered,
                     party5=party5.reordered,
                     party6=party6.reordered)


mydata$partisanship <- NA
mydata$partisanship[mydata$party2 == 1] <- "Lean Democrat"
mydata$partisanship[mydata$party2 == 2] <- "Independent"
mydata$partisanship[mydata$party2 == 3] <- "Lean Republican"
mydata$partisanship[mydata$party2 == 4] <- "Not sure"

mydata$partisanship[mydata$party3 == 1] <- "Lean Democrat"
mydata$partisanship[mydata$party3 == 2] <- "Independent"
mydata$partisanship[mydata$party3 == 3] <- "Lean Republican"
mydata$partisanship[mydata$party3 == 4] <- "Not sure"

mydata$partisanship[mydata$party4 == 1] <- "Lean Democrat"
mydata$partisanship[mydata$party4 == 2] <- "Independent"
mydata$partisanship[mydata$party4 == 3] <- "Lean Republican"
mydata$partisanship[mydata$party4 == 4] <- "Not sure"

mydata$partisanship[mydata$party5 == 1] <- "Strong Democrat"
mydata$partisanship[mydata$party5 == 2] <- "Not very strong Democrat"

mydata$partisanship[mydata$party6 == 1] <- "Strong Republican"
mydata$partisanship[mydata$party6 == 2] <- "Not very strong Republican"

print(table(mydata$partisanship))





## ideology, partisanship, news 

pdf(file="Partisanship-Ideology-News-gamma-newIDs-AttnCheckPassed.pdf", height=3.5, width=13)
    layout(mat=matrix(c(1,2,3,4,5,6), 2, 3), heights=c(5,2))
for (vname in c("strong.repub", "ideology", "news.lean")){
    formula.string <- paste(vname, "~ lp(gamma.hat, nn=0.7, deg=1)")

    n.samp <- M
    sample.inds <- sample(1:M, size=n.samp, replace=FALSE)
    vals <- seq(from=min(gamma.hat), to=max(gamma.hat), length.out=500)
    prediction.mat <- matrix(NA, n.samp, length(vals))
    for (i in 1:n.samp){
        if (i %% 50 == 0){
            cat("sample", i, "of", n.samp, "\n")
        }
        mydata.samp <- mydata
        mydata.samp$gamma.hat <- gammas[sample.inds[i],]
        loc.out <- locfit(as.formula(formula.string), data=mydata.samp)
        prediction.mat[i,] <- predict(loc.out, data.frame(gamma.hat=vals))      
    }
    lower.pred <- apply(prediction.mat, 2, quantile, prob=0.025)
    upper.pred <- apply(prediction.mat, 2, quantile, prob=0.975)
    median.pred <- apply(prediction.mat, 2, quantile, prob=0.5)

    

    ylim.v <- c(0,1)
    hline.vals <- seq(from=0, to=1, by=0.2)
    yname <- "Fraction Strong Republican"
    if (vname == "ideology"){
        ylim.v <- c(1,7)
        hline.vals <- seq(from=0, to=10, by=1)
        yname <- "Ideology"
    }

    if (vname == "news.lean"){
        ylim.v <- c(-2,2)
        hline.vals <- seq(from=-10, to=10, by=1)
        yname <- "Slant of News Consumption"
    }
    
    
    
    par(mar=c(2,4,1,2), lend=1)
    plot(vals, median.pred,
         type="l", xlim=c(0, pi/2), ylim=ylim.v, xlab="", ylab=yname, lwd=2,
         col=rgb(1, 0.169, 0, 1))
    abline(v=seq(from=0, to=1.5, by=0.25), col="lightgray")
    abline(h=hline.vals, col="lightgray")
    lines(vals, median.pred, col=rgb(1, 0.169, 0, 1), lwd=2)
    polygon(x=c(vals, rev(vals)), y=c(lower.pred, rev(upper.pred)),
            col=rgb(1, 0.169, 0, 0.3), border=NA)
    par(mar=c(4.5,4,0,2))
    hist(gamma.hat, nclass=50, col=rgb(0.6, 0.933, 1, 1),
         prob=TRUE, ylim=c(0,3.5),
         main="", xlab=expression(paste("E[", gamma[i], "|Y]", sep="")),
         xlim=c(0, pi/2))

    
}
dev.off()




















## nosocdist, nmask 
pdf(file="SocialDistancing-MaskWearing-gamma-newIDs-AttnCheckPassed.pdf", height=4.5, width=11)
    layout(mat=matrix(c(1,2,3,4), 2, 2), heights=c(5,2))
for (vname in c("nosocdist", "nmask")){
    formula.string <- paste(vname, "~ lp(gamma.hat, nn=0.7, deg=1)")

    n.samp <- M
    sample.inds <- sample(1:M, size=n.samp, replace=FALSE)
    vals <- seq(from=min(gamma.hat), to=max(gamma.hat), length.out=500)
    prediction.mat <- matrix(NA, n.samp, length(vals))
    for (i in 1:n.samp){
        if (i %% 50 == 0){
            cat("sample", i, "of", n.samp, "\n")
        }
        mydata.samp <- mydata
        mydata.samp$gamma.hat <- gammas[sample.inds[i],]
        loc.out <- locfit(as.formula(formula.string), data=mydata.samp)
        prediction.mat[i,] <- predict(loc.out, data.frame(gamma.hat=vals))      
    }
    lower.pred <- apply(prediction.mat, 2, quantile, prob=0.025)
    upper.pred <- apply(prediction.mat, 2, quantile, prob=0.975)
    median.pred <- apply(prediction.mat, 2, quantile, prob=0.5)

    

    ylim.v <- c(0,1)
    hline.vals <- seq(from=0, to=1, by=0.2)
    yname <- "Fraction Not Social Distancing"
    if (vname == "nmask"){
        ylim.v <- c(0,9)
        hline.vals <- seq(from=0, to=10, by=2)
        yname <- "Amount of Mask Wearing"
    }

    
    
    par(mar=c(2,4,1,2), lend=1)
    plot(vals, median.pred,
         type="l", xlim=c(0, pi/2), ylim=ylim.v, xlab="", ylab=yname, lwd=2,
         col=rgb(1, 0.169, 0, 1))
    abline(v=seq(from=0, to=1.5, by=0.25), col="lightgray")
    abline(h=hline.vals, col="lightgray")
    lines(vals, median.pred, col=rgb(1, 0.169, 0, 1), lwd=2)
    polygon(x=c(vals, rev(vals)), y=c(lower.pred, rev(upper.pred)),
            col=rgb(1, 0.169, 0, 0.3), border=NA)
    par(mar=c(4.5,4,0,2))
    hist(gamma.hat, nclass=50, col=rgb(0.6, 0.933, 1, 1),
         prob=TRUE, ylim=c(0,3.5),
         main="", xlab=expression(paste("E[", gamma[i], "|Y]", sep="")),
         xlim=c(0, pi/2))

    
}
dev.off()






